load("//scann/oss_wrappers:build_defs.bzl", "scann_py_proto_library")

# -*- mode: python; -*-
#
package(
    default_visibility = ["//visibility:public"],
    licenses = ["notice"],
)

# Protos
# =========================================================

proto_library(
    name = "kmeans_tree_proto",
    srcs = ["kmeans_tree.proto"],
    tags = ["local"],
    deps = ["//scann/proto:partitioning_proto"],
)

cc_proto_library(
    name = "kmeans_tree_cc_proto",
    tags = ["local"],
    deps = [":kmeans_tree_proto"],
)

scann_py_proto_library(
    name = "kmeans_tree_py_pb2",
    srcs = ["kmeans_tree.proto"],
    proto_deps = [":kmeans_tree_proto"],
    py_proto_deps = [
        "//scann/proto:partitioning_py_pb2",
    ],
    tags = ["local"],
)

# Library
# =========================================================
cc_library(
    name = "kmeans_tree",
    srcs = [
        "kmeans_tree.cc",
        "kmeans_tree_node.cc",
        "training_options.cc",
    ],
    hdrs = [
        "kmeans_tree.h",
        "kmeans_tree_node.h",
        "training_options.h",
    ],
    tags = ["local"],
    deps = [
        ":kmeans_tree_cc_proto",
        "//scann/data_format:datapoint",
        "//scann/data_format:dataset",
        "//scann/distance_measures:distance_measure_base",
        "//scann/distance_measures/one_to_many",
        "//scann/distance_measures/one_to_one:l2_distance",
        "//scann/oss_wrappers:scann_aligned_malloc",
        "//scann/oss_wrappers:scann_castops",
        "//scann/oss_wrappers:scann_down_cast",
        "//scann/oss_wrappers:scann_random",
        "//scann/oss_wrappers:scann_status",
        "//scann/oss_wrappers:scann_threadpool",
        "//scann/oss_wrappers:tf_dependency",
        "//scann/partitioning:anisotropic",
        "//scann/proto:partitioning_cc_proto",
        "//scann/utils:common",
        "//scann/utils:datapoint_utils",
        "//scann/utils:fast_top_neighbors",
        "//scann/utils:gmm_utils",
        "//scann/utils:scalar_quantization_helpers",
        "//scann/utils:types",
        "//scann/utils:util_functions",
        "//scann/utils:zip_sort",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/flags:flag",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/random",
        "@com_google_absl//absl/strings",
    ],
)
